// src/context/index.ts
import { create } from "zustand";
import { persist } from "zustand/middleware";
import { Agent, Message, ChatHistoryRecord, ChatMessage } from "../types";

interface AgentState {
  // 当前选中的 Agent
  selectedAgent: Agent | null;

  // 设置选中的 Agent
  setSelectedAgent: (agent: Agent) => void;

  // 清除选中的 Agent
  clearSelectedAgent: () => void;
}

// 创建 Agent 状态管理
export const useAgentStore = create<AgentState>()(
  persist(
    (set) => ({
      selectedAgent: null,

      setSelectedAgent: (agent: Agent) => set({ selectedAgent: agent }),

      clearSelectedAgent: () => set({ selectedAgent: null }),
    }),
    {
      name: "agent-storage", // 本地存储的键名
      partialize: (state) => ({ selectedAgent: state.selectedAgent }), // 只持久化 selectedAgent
    }
  )
);

// 聊天历史记录状态管理
interface ChatHistoryState {
  // 当前选中的历史记录
  selectedHistory: ChatHistoryRecord | null;

  // 设置选中的历史记录
  setSelectedHistory: (history: ChatHistoryRecord) => void;

  // 清除选中的历史记录
  clearSelectedHistory: () => void;

  // 向选中的历史记录添加新消息
  addMessageToHistory: (message: Message) => void;

  // 更新选中的历史记录中的消息
  updateHistoryMessages: (messages: Message[]) => void;
}

// 创建聊天历史记录状态管理
export const useChatHistoryStore = create<ChatHistoryState>()(
  persist(
    (set) => ({
      selectedHistory: null,

      setSelectedHistory: (history: ChatHistoryRecord) =>
        set({ selectedHistory: history }),

      clearSelectedHistory: () => set({ selectedHistory: null }),

      addMessageToHistory: (message: Message) =>
        set((state) => {
          if (!state.selectedHistory) return state;
          
          // 确保 history 是数组
          const historyArray = Array.isArray(state.selectedHistory.history) 
            ? state.selectedHistory.history 
            : [];
            
          // 将 Message 转换为 ChatMessage
          const chatMessage: ChatMessage = {
            role: message.isUser ? "user" : "assistant",
            message: message.content
          };

          return {
            selectedHistory: {
              ...state.selectedHistory,
              history: [...historyArray, chatMessage],
            },
          };
        }),

      updateHistoryMessages: (messages: Message[]) =>
        set((state) => {
          if (!state.selectedHistory) return state;

          // 将 Message[] 转换为 ChatMessage[]
          const chatMessages: ChatMessage[] = messages.map(msg => ({
            role: msg.isUser ? "user" : "assistant",
            message: msg.content
          }));

          return {
            selectedHistory: {
              ...state.selectedHistory,
              history: chatMessages,
            },
          };
        }),
    }),
    {
      name: "chat-history-storage", // 本地存储的键名
      partialize: (state) => ({ selectedHistory: state.selectedHistory }), // 只持久化 selectedHistory
    }
  )
);

// 用户状态管理
export interface User {
  username: string; // 手机号
  nickname?: string; // 昵称
  VIPExpire?: string;
  isLoggedIn: boolean;
}

interface UserState {
  user: User | null;
  setUser: (user: User) => void;
  clearUser: () => void;
  updateUserInfo: (userInfo: Partial<User>) => void;
}

// 创建用户状态管理
export const useUserStore = create<UserState>()(
  persist(
    (set) => ({
      user: null,

      setUser: (user: User) => set({ user }),

      clearUser: () => set({ user: null }),

      updateUserInfo: (userInfo: Partial<User>) =>
        set((state) => ({
          user: state.user ? { ...state.user, ...userInfo } : null,
        })),
    }),
    {
      name: "user-storage", // 本地存储的键名
      partialize: (state) => ({ user: state.user }), // 只持久化 user
    }
  )
);
